/*    file: AVFrameUtil.cpp
 *    desc:
 * 
 * created: 2016-06-12
 *  author: chuanjiang.zh@qq.com
 * company: 
 */

#include "BasicType.h"
#include "AVFrameUtil.h"
#include "TStringUtil.h"

#include <stdio.h>
#include <stdlib.h>

#ifdef WIN32

#else

#pragma pack(1)

typedef unsigned char BYTE;
typedef unsigned short WORD;
typedef unsigned int   DWORD;
typedef int	LONG;

typedef struct tagRGBQUAD {
    BYTE rgbBlue;
    BYTE rgbGreen;
    BYTE rgbRed;
    BYTE rgbReserved;
} RGBQUAD;

typedef struct tagBITMAPFILEHEADER {
    WORD bfType;
    DWORD bfSize;
    WORD bfReserved1;
    WORD bfReserved2;
    DWORD bfOffBits;
} BITMAPFILEHEADER,*LPBITMAPFILEHEADER,*PBITMAPFILEHEADER;

typedef struct tagBITMAPINFOHEADER {
    DWORD biSize;
    LONG biWidth;
    LONG biHeight;
    WORD biPlanes;
    WORD biBitCount;
    DWORD biCompression;
    DWORD biSizeImage;
    LONG biXPelsPerMeter;
    LONG biYPelsPerMeter;
    DWORD biClrUsed;
    DWORD biClrImportant;
} BITMAPINFOHEADER,*LPBITMAPINFOHEADER,*PBITMAPINFOHEADER;

#pragma()

#define BI_RGB 	0

#endif //


extern "C"
{
#include <libswscale/swscale.h>
}


static bool yuv420pToJpeg(AVFrame *pFrame, const char* filename)
{
    AVFormatContext* pFormatCtx = avformat_alloc_context();
    pFormatCtx->oformat = av_guess_format("mjpeg", NULL, NULL);
    if (!pFormatCtx->oformat)
    {
        avformat_free_context(pFormatCtx);
        return false;
    }

    if (avio_open(&pFormatCtx->pb, filename, AVIO_FLAG_READ_WRITE) < 0)
    {
        avformat_free_context(pFormatCtx);
        return false;
    }

    AVStream* video_st = avformat_new_stream(pFormatCtx, 0);
    if (video_st == NULL)
    {
        avformat_free_context(pFormatCtx);
        return false;
    }

    AVCodecContext* pCodecCtx = video_st->codec;
    pCodecCtx->codec_id = pFormatCtx->oformat->video_codec;
    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;

    pCodecCtx->width = pFrame->width;
    pCodecCtx->height = pFrame->height;

    video_st->time_base.num = 1;
    video_st->time_base.den = 25;

    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = 25;

    AVCodec *jpegCodec = avcodec_find_encoder(pCodecCtx->codec_id);
    if (!jpegCodec) {
        return false;
    }

    if (avcodec_open2(pCodecCtx, jpegCodec, NULL) < 0) 
    {
        avformat_free_context(pFormatCtx);
        return false;
    }

    avformat_write_header(pFormatCtx, NULL);


    AVPacket packet;
    av_init_packet(&packet);
    packet.data = NULL;
    packet.size = 0;

    int gotFrame;

    avcodec_encode_video2(pCodecCtx, &packet, pFrame, &gotFrame);
    if (gotFrame)
    {
        av_write_frame(pFormatCtx, &packet);
    }

    av_packet_unref(&packet);

    av_write_trailer(pFormatCtx);

    if (video_st)
    {
        avcodec_close(video_st->codec);
    }

    avio_close(pFormatCtx->pb);
    avformat_free_context(pFormatCtx);

    return true;
}

bool AVFrameUtil::saveAsJpeg(AVFrame *pFrame, const char* filename)
{
	bool done = false;
    AVFrame* rgbFrame = pFrame;
    if (pFrame->format != AV_PIX_FMT_YUV420P)
    {
        AVFrame* rgbFrame = av_frame_alloc();
        rgbFrame->width = pFrame->width;
        rgbFrame->height = pFrame->height;
        rgbFrame->format = AV_PIX_FMT_YUV420P;
        rgbFrame->pts = pFrame->pts;
        av_frame_get_buffer(rgbFrame, 1);

        SwsContext* pContext = sws_getContext(pFrame->width,
            pFrame->height, (AVPixelFormat)pFrame->format,
            pFrame->width, pFrame->height,
            AV_PIX_FMT_YUV420P, SWS_BILINEAR, 0, 0, 0);
		if (pContext)
		{
			sws_scale(pContext, pFrame->data, pFrame->linesize, 0, pFrame->height,
				rgbFrame->data, rgbFrame->linesize);

			done = yuv420pToJpeg(rgbFrame, filename);

			sws_freeContext(pContext);
		}
        av_frame_free(&rgbFrame);
        return done;
    }
    return yuv420pToJpeg(pFrame, filename);
}



static bool rgb32FrameToBitmap(AVFrame* frame, const char* filename)
{
#define BFT_BITMAP 0x4d42   /* 'BM' */
#define DibNumColors(lpbi)      ((lpbi)->biClrUsed == 0 && (lpbi)->biBitCount <= 8 \
    ? (int)(1 << (int)(lpbi)->biBitCount)          \
    : (int)(lpbi)->biClrUsed)

#define DibSize(lpbi)           ((lpbi)->biSize + (lpbi)->biSizeImage + (int)(lpbi)->biClrUsed * sizeof(RGBQUAD))

#define DibPaletteSize(lpbi)    (DibNumColors(lpbi) * sizeof(RGBQUAD))

    BITMAPFILEHEADER fileHeader;
    memset(&fileHeader, 0, sizeof(fileHeader));

    BITMAPINFOHEADER infoHeader;
    memset(&infoHeader, 0, sizeof(infoHeader));
    infoHeader.biSize = sizeof(infoHeader);
    infoHeader.biWidth = frame->width;
    infoHeader.biHeight = -frame->height;
    infoHeader.biPlanes = 1;
    infoHeader.biBitCount = 32;
    infoHeader.biCompression = BI_RGB;

    LPBITMAPINFOHEADER  pdib = (LPBITMAPINFOHEADER)&infoHeader;

    DWORD dwSize = DibSize(pdib);
    fileHeader.bfType          = BFT_BITMAP;
    fileHeader.bfSize          = dwSize + sizeof(BITMAPFILEHEADER);
    fileHeader.bfReserved1     = 0;
    fileHeader.bfReserved2     = 0;
    fileHeader.bfOffBits       = (DWORD)sizeof(BITMAPFILEHEADER) + pdib->biSize + DibPaletteSize(pdib);

    bool done = false;
    FILE* file = NULL;
    file = fopen(filename, "wb");
    if (file)
    {
        fwrite(&fileHeader, 1, sizeof(BITMAPFILEHEADER), file);
        fwrite(pdib, 1, infoHeader.biSize, file);
        fwrite(frame->data[0], 1, frame->linesize[0] * frame->height, file);
        fclose(file);
        done = true;
    }
    return done;
}

bool AVFrameUtil::saveAsBmp(AVFrame *frame, const char* filename)
{
    if (!frame || !frame->width)
    {
        return false;
    }

    if (frame->format != AV_PIX_FMT_BGRA)
    {
        AVFrame* rgbFrame = av_frame_alloc();
        rgbFrame->width = frame->width;
        rgbFrame->height = frame->height;
        rgbFrame->format = AV_PIX_FMT_BGRA;
        av_frame_get_buffer(rgbFrame, 1);

        SwsContext* pContext = sws_getContext(frame->width,
            frame->height, (AVPixelFormat)frame->format,
            frame->width, frame->height,
            AV_PIX_FMT_BGRA, SWS_BILINEAR, 0, 0, 0);

        sws_scale(pContext, frame->data, frame->linesize, 0, frame->height,
            rgbFrame->data, rgbFrame->linesize);

        bool done = rgb32FrameToBitmap(rgbFrame, filename);

        sws_freeContext(pContext);

        av_frame_free(&rgbFrame);
        return done;
    }
    else
    {
        return rgb32FrameToBitmap(frame, filename);
    }
}

bool AVFrameUtil::saveFrame(AVFrame *pFrame, const char* filename)
{
    std::string head;
    std::string ext;
    comn::StringUtil::rsplit(filename, '.', head, ext);   
    comn::StringUtil::toLower(ext);
    if ((ext == "jpeg") || (ext == "jpg"))
    {
        return saveAsJpeg(pFrame, filename);
    }
    else
    {
        return saveAsBmp(pFrame, filename);
    }
}